rm(list = ls())
library(tidyverse)
library(estimatr)
library(texreg)
library(ggrepel)
library(gridExtra)
library(ggridges)
library(ggplot2)
library(ggpubr)
library(grid)
setwd(dirname(rstudioapi::getActiveDocumentContext()$path))
source('fn.R')

load('data.rdata')

mod <- lm_robust(Y ~ arm - 1, data = data, weights = ipw)

gg_df <- tidy(mod) %>%
  mutate(label = gsub(pattern = 'arm', replacement = 'Arm ', x = term),
         label = fct_reorder(factor(label), estimate)) %>%
  left_join(data %>% group_by(arm) %>% summarize(arm_n = n()) %>% mutate(term = paste0('arm', arm))) %>%
  mutate(e_label = paste0(format_num(estimate, 3), " ", add_parens(std.error, 3)),
         t_label = ifelse(term == 'arm10',
                          paste0('Est: ', format_num(estimate, 3), ' (SE: ', format_num(std.error, 3), ') [N = ', arm_n, ']'),
                                 paste0(e_label, ' [', arm_n, ']')))

p_est <- ggplot(gg_df, aes(x = estimate, y = label, xmin = conf.low, xmax = conf.high)) +
  geom_vline(xintercept = 0, linetype = 'dotted') +
  geom_point() + geom_errorbarh(height = 0) + 
  xlab('Average proportion of respondents supporting the measure') +
  ylab('Treatment Arm') +
  theme_bw() + 
  geom_text(aes(label = t_label), nudge_y = .3, nudge_x = 0, size = 2.5, show.legend = FALSE) +
  scale_color_manual(values = 'black') + 
  theme(strip.background = element_blank(),
        axis.title.y = element_blank(),
        legend.title = element_blank(),
        legend.position = 'bottom',
        plot.title = element_text(hjust = 0.5)) +
  xlim(0, 1) 

ggsave(file="Figures/Figure 1.pdf", p_est, width = 5, height = 4)
